import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.statespace.sarimax import SARIMAX
from statsmodels.graphics.gofplots import qqplot
from statsmodels.tools.sm_exceptions import ConvergenceWarning
import warnings


# 数据处理函数（沿用之前的）
def process_data():
    excel_file = pd.ExcelFile("ori_message.xlsx")
    years = ['2016', '2017', '2018', '2019', '2020', '2021']
    all_monthly_data = []
    for year in years:
        df = excel_file.parse(year)
        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3) ': 'mean'
        }).reset_index()
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3) ']
        grouped = grouped.round({
            '流量(m3/s)': 2,
            '含沙量(kg/m3) ': 4,
            '排沙量(kg/s)': 2
        })
        all_monthly_data.append(grouped)
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result = result.rename(columns={'含沙量(kg/m3) ': '含沙量(kg/m3)'})
    result['时间'] = pd.to_datetime(result['年'].astype(str) + '-' + result['月'].astype(str) + '-01')
    return result


if __name__ == "__main__":
    processed_data = process_data()
    processed_data.set_index('时间', inplace=True)  # 设置时间为索引
    flow_series = processed_data['流量(m3/s)']

    # 忽略收敛警告
    warnings.filterwarnings("ignore", category=ConvergenceWarning)

    # 确定SARIMA模型参数
    p, d, q = 2, 2, 3
    P, D, Q, S = 0, 1, 1, 12  # 假设季节性周期S=12（月度数据）

    # 拟合SARIMA模型
    model = SARIMAX(flow_series, order=(p, d, q), seasonal_order=(P, D, Q, S))
    results = model.fit()

    # 获取残差
    residuals = results.resid

    # 绘制QQ图
    plt.figure(figsize=(8, 6))
    qqplot(residuals, line='s', ax=plt.gca())
    plt.title('图26.贴合度图')
    plt.xlabel('Theoretical Quantiles')
    plt.ylabel('Sample Quantiles')
    plt.tight_layout()
    plt.savefig('图26.贴合度图.png', dpi=300, bbox_inches='tight')
    plt.show()